import torch

from pointnet2_with_pcld_condition import PointNet2CloudCondition

from json_reader import restore_string_to_list_in_a_dict

import json


def main():
    points = torch.randn(1, 2048, 3).cuda()
    condition = torch.randn(1, 3072, 4).cuda()
    point2_config_path = "/home/zhangtaihao/workspace/code/diff/pointNet-demo/config/config.json"
    with open(point2_config_path) as f:
        data = f.read()
    config = json.loads(data)
    config = restore_string_to_list_in_a_dict(config)
    pointnet_config = config["pointnet_config"]

    net = PointNet2CloudCondition(pointnet_config).cuda()
    t = torch.randint(1000, size=(1,)).cuda()
    label = torch.tensor([1]).cuda()
    output = net(points, condition, t, label=label)
    print(output.shape)

    pass


if __name__ == "__main__":
    main()
